# Author: Olivier Grisel <olivier.grisel@ensta.org>
#         Lars Buitinck <L.J.Buitinck@uva.nl>
#         Chyi-Kwei Yau <chyikwei.yau@gmail.com>
# License: BSD 3 clause

from __future__ import print_function
from time import time

from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from sklearn.decomposition import NMF, LatentDirichletAllocation
from sklearn.datasets import fetch_20newsgroups

# n_samples = 2000
# n_features = 1000
n_topics = 10
n_top_words = 20

root = "G:/experiment/revert_lda/news/"


def write_phi(data, wordmap):
    with open(root+"model-final.phi", 'w', encoding='utf-8') as w:
        for line in data:
            w.write(' '.join([str(item) for item in line])+"\n")

    with open(root + "model-final.wordmap", 'w', encoding='utf-8') as w:
        for key in range(len(wordmap)):
            w.write(wordmap[key] +" "+str(key) + "\n")

def read_data(path):
    result = None
    with open(path, 'r',encoding='utf-8') as f:
        result = f.readlines()

    return result


def print_top_words(model, feature_names, n_top_words):
    for topic_idx, topic in enumerate(model.components_):
        print("Topic #%d:" % topic_idx)
        print(" ".join([feature_names[i]
                        for i in topic.argsort()[:-n_top_words - 1:-1]]))
    print()


# Load the 20 newsgroups dataset and vectorize it. We use a few heuristics
# to filter out useless terms early on: the posts are stripped of headers,
# footers and quoted replies, and common English words, words occurring in
# only one document or in at least 95% of the documents are removed.

print("Loading dataset...")
t0 = time()
# dataset = fetch_20newsgroups(shuffle=True, random_state=1,
#                              remove=('headers', 'footers', 'quotes'))

path = root+"news.data.worddoc"
data_samples = read_data(path)
print("done in %0.3fs." % (time() - t0))

# Use tf-idf features for NMF.
print("Extracting tf-idf features for NMF...")
tfidf_vectorizer = TfidfVectorizer()
t0 = time()
print(len(data_samples), data_samples[0])
tfidf = tfidf_vectorizer.fit_transform(data_samples)
print("done in %0.3fs." % (time() - t0))

# Use tf (raw term count) features for LDA.
# print("Extracting tf features for LDA...")
# tf_vectorizer = CountVectorizer(max_df=0.95, min_df=2, max_features=n_features,
#                                 stop_words='english')
# t0 = time()
# tf = tf_vectorizer.fit_transform(data_samples)
# print("done in %0.3fs." % (time() - t0))

# Fit the NMF model
# print("Fitting the NMF model with tf-idf features,"
#       "n_samples=%d and n_features=%d..."
#       % (n_samples, n_features))
t0 = time()
nmf = NMF(n_components=n_topics, max_iter=2000, nls_max_iter=2000).fit(tfidf)
tfidf_feature_names = tfidf_vectorizer.get_feature_names()
write_phi(nmf.components_, tfidf_feature_names)
# exit()
print("done in %0.3fs." % (time() - t0))

print("\nTopics in NMF model:")

print_top_words(nmf, tfidf_feature_names, n_top_words)
#
# print("Fitting LDA models with tf features, n_samples=%d and n_features=%d..."
#       % (n_samples, n_features))
# lda = LatentDirichletAllocation(n_topics=n_topics, max_iter=5,
#                                 learning_method='online', learning_offset=50.,
#                                 random_state=0)
# t0 = time()
# print(tfidf)
# lda.fit(tfidf)
# print("done in %0.3fs." % (time() - t0))
#
# print("\nTopics in LDA model:")
# tf_feature_names = tf_vectorizer.get_feature_names()
# print_top_words(lda, tfidf_feature_names, n_top_words)